package rate_limit

import (
	"net/http"
	"time"

	"github.com/didip/tollbooth/v6"
	"github.com/didip/tollbooth/v6/limiter"
	"github.com/gin-gonic/gin"

	"github.com/super-team/start-kit/configs"
)

func RateLimit() gin.HandlerFunc {
	cfg := configs.Config.RateLimit
	l := tollbooth.NewLimiter(cfg.Max, &limiter.ExpirableOptions{
		DefaultExpirationTTL: time.Duration(cfg.Expire) * time.Second,
	})
	return func(ctx *gin.Context) {
		if err := tollbooth.LimitByRequest(l, ctx.Writer, ctx.Request); err != nil {
			ctx.JSON(http.StatusTooManyRequests, gin.H{
				"code": http.StatusTooManyRequests,
				"msg":  "请求太快了，慢点吧",
				"data": nil,
			})
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}
